import ctypes
from ctypes import *
import time
import sys
import threading
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
import xml.etree.ElementTree as ET
import os
import sys
from PyQt5.QtGui import QPalette, QPixmap, QColor, QPainter, \
    QPen,QBrush,QPainterPath,QPolygon,QFont,QIcon,QImageReader
# from PyQt5.Qt import QImageReader
from PIL import Image
from VOCdefects.Detector_Struct import *
from universal.root_path import dir_root
from mylibs.CompressionImages import Compress_Picture
FILE_PATH = os.path.dirname(sys.argv[0])
import matplotlib.pyplot as plt

class trainValue():
    def __init__(self):
        self.current_batch = 0
        self.avg_loss = 0
        self.current_loss = 0
        self.batchTime = 0
        self.weightName = ''
        self.trianTime = ''

class runTrain(QThread):
    def __init__(self,projectPath,trainImgspath,labelAnnatationPath,preWeightFile,
                 saveWeigthPath,batchsize,inputImageSize,epochs,
                 learningRate,classesList,trainPeriod,inputWeightName,finalName):
        super().__init__()
        self.dirpath = dir_root()
        self.dllPath = os.sep+"lib"+os.sep
        self.dllPath2 = os.sep+"TrainData"+os.sep
        imageSize = inputImageSize.split('*')
        self.trainImgsPath = trainImgspath
        self.labelAnnatationPath = labelAnnatationPath
        self.proDir = os.path.dirname(projectPath)
        self.proName = os.path.basename(projectPath)
        self.saveWeigthPath = saveWeigthPath
        self.classes = classesList
        self.trainParameters = trainParameters()
        self.result_paramsTmp = result_params()  #
        self.strcfg = bytes(self.dirpath+self.dllPath2 + "yolov4-tiny.cfg", 'utf-8')
        print(self.strcfg,'self.strcfg')
        self.dataCfg = bytes(self.dirpath+self.dllPath2   + "obj.data", 'utf-8')
        print(self.dataCfg,'self.dataCfg')
        if preWeightFile==None or preWeightFile.strip()=='':
            self.weightfile = bytes(self.dirpath+self.dllPath2  + "yolov4-tiny.conv.29", 'utf-8')
            print(self.weightfile,'self.weightfile')
        else:
            self.weightfile = bytes(preWeightFile,'utf-8')#
        self.iTrain = ctypes.c_int(0)
        self.gpus = c_int(0)
        self.thread_train = threading.Thread(target=self.train,args='')
        self.thread_train.daemon = True
        # self.thread_getValue = threading.Thread( target=self.getValue,args='')
        # self.thread_getValue.daemon = True
        # self.thread_getValue.start()
        self.allValues =[]
        self.trainParameters.classes = len(classesList)
        self.setObjParam()
        self.trainvalue = trainValue
        self.current_batch = 0
        self.avg_loss = 0
        self.current_loss = 0
        self.batchTime = 0
        self.weightName = ''

        # self.trainParameters.batch = batchsize
        # self.trainParameters.subdivisions = 16

        self.trainParameters.batch = batchsize
        self.trainParameters.subdivisions = 4

        self.trainPeriod = trainPeriod
        print(self.trainParameters.subdivisions)
        self.trainParameters.width = int(imageSize[0])
        self.trainParameters.height = int(imageSize[1])
        self.trainParameters.max_batches = epochs
        self.trainParameters.learning_rate = learningRate
        self.trainParameters.filters = (self.trainParameters.classes + 5)*3
        self.trainParameters.inPutWeightName = bytes(inputWeightName, 'utf-8')
        self.trainParameters.inPutFinalWeightName = bytes(finalName, 'utf-8')
        self.voc_generation()
        plt.ion()

    def run(self):
        print('run')
        self.thread_train.start()
        # return self.getValue()

    def setObjParam(self):
        with open(self.dirpath+self.dllPath2  + "obj.data",'w+') as objF:
            objF.write("classes={}".format(self.trainParameters.classes)+'\n')
            objF.write("train = " + self.dirpath+self.dllPath2   + "train.txt"+'\n')
            objF.write("test = " + self.dirpath+self.dllPath2  + "test.txt"+'\n')
            # objF.write("backup = " +self.proDir + '/')
            objF.write("backup = " + self.saveWeigthPath + '/')

        with open(self.dirpath+self.dllPath2 +"obj.names",'w') as objName:
            for cls in self.classes:
                objName.write(cls+'\n')

    def voc_generation(self):
        vocLabel = generateLabel_voc(self.trainImgsPath,self.classes,self.labelAnnatationPath)
        vocLabel.startVoc()

    def printValues(self):
        return self.allValues

    def train(self):
        print('self.Objdll',self.dirpath+self.dllPath+"Detector_API.dll")
        self.Objdll = ctypes.windll.LoadLibrary(r'%s'%(self.dirpath+self.dllPath+"Detector_API.dll"))
        self.Objdll.InitDetectorDll.restype = c_bool
        self.Objdll.InitDetectorDll()
        beginTrain = time.time()#
        self.Objdll.train_detector_img.restype = c_voidp
        self.Objdll.train_detector_img(ctypes.byref(self.trainParameters), self.dataCfg,
                                    self.strcfg, self.weightfile,
                                    ctypes.byref(self.gpus),ctypes.c_int(1),
                                    ctypes.c_int(self.trainPeriod), ctypes.byref(self.result_paramsTmp))#

        self.current_batch = str(self.result_paramsTmp.current_batch)
        self.trainvalue.current_batch = str(self.result_paramsTmp.current_batch)
        self.trainvalue.avg_loss = str('%.2f'%self.result_paramsTmp.dAvg_loss)
        self.trainvalue.current_loss = str('%.2f'%self.result_paramsTmp.current_loss)
        #======================================================
        self.trainvalue.batchTime = str(self.result_paramsTmp.dTime)
        endTrain = time.time()
        self.trainvalue.trianTime =str('%.3f'%(endTrain - beginTrain))
        # self.Objdll.FreeDetectorDll.restype = c_bool
        # self.Objdll.FreeDetectorDll()
        return self.trainvalue

    def sleeptime(self,hour,min,sec):
        return hour*3600 + min * 60 + sec


    def getValue(self):
        while True:
            time.sleep(0.1)
            if (self.current_batch != self.result_paramsTmp.current_batch):
                self.current_batch = self.result_paramsTmp.current_batch
                self.trainvalue.current_batch = self.result_paramsTmp.current_batch
                self.trainvalue.avg_loss = self.result_paramsTmp.dAvg_loss
                self.trainvalue.current_loss = self.result_paramsTmp.current_loss
                #========================================
                print(self.trainvalue.current_loss)
                self.trainvalue.batchTime = self.result_paramsTmp.dTime
                self.trainvalue.weightName = self.result_paramsTmp.weightName.decode()
                self.allValues.append(self.trainvalue)
            if (self.current_batch == self.trainParameters.max_batches):
                self.printValues()
                break

class generateLabel_voc():
    def __init__(self,trainImgsPath,classes,annatationPath):
        self.dirpath = dir_root()
        self.dllPath = os.sep + "lib" + os.sep
        self.dllPath2 = os.sep + "TrainData" + os.sep
        self.trainImgsPath = trainImgsPath
        self.classes = classes
        self.annatationPath = annatationPath
        self.dirPath = FILE_PATH + '/'
        self.imageType = ''

    def startVoc(self):
        self.generateValid()
        image_ids = open(self.dirpath+self.dllPath2+'valid.txt').read().strip().split()
        list_file = open(self.dirpath+self.dllPath2+'train.txt', 'w+')
        list_file.truncate()
        for image_id in image_ids:
            list_file.write(self.trainImgsPath+'/'+ '%s.%s\n' % (image_id,self.imageType))
            self.convert_annotation(image_id)
        list_file.close()

    def generateValid(self):
        with open(self.dirpath+self.dllPath2 + "valid.txt", 'w') as f:
            f.truncate()
            print(os.walk(self.trainImgsPath))
            for filenames in os.walk(self.trainImgsPath):
                filenames = list(filenames)
                filenames = filenames[2]
            for filename in filenames:
                a = filename.split('.')
                extensions = ['.%s' % fmt.data().decode("ascii").lower() for fmt in
                              QImageReader.supportedImageFormats()]
                if '.'+a[1] in extensions:
                    f.write(a[0] + '\n')
                    self.imageType = a[1]
        return

    def convert(self,size, box):
        dw = 1./(size[0])
        dh = 1./(size[1])
        x = (box[0] + box[1])/2.0
        y = (box[2] + box[3])/2.0
        w = box[1] - box[0]
        h = box[3] - box[2]
        x = x*dw
        w = w*dw
        y = y*dh
        h = h*dh
        return (x,y,w,h)

    def convert_annotation(self,image_id):
        in_file = open(self.annatationPath+'/'+'%s.xml'%(image_id))
        out_file = open(self.annatationPath+'/'+'%s.txt'%(image_id), 'w')
        tree=ET.parse(in_file)
        root = tree.getroot()
        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)
        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            cls = obj.find('name').text
            if cls not in self.classes or int(difficult)==1:
                continue
            cls_id = self.classes.index(cls)
            xmlbox = obj.find('bndbox')
            b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
            if w ==0 or h==0:
                w=2448
                h=2048
            bb = self.convert((w,h), b)
            out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')










